import importlib.util
import os
import numpy as np
from ml_dtypes import bfloat16


def compute(case_desc, timestamp):
    expect_func = case_desc.expect_func

    # 假设 b.py 的绝对路径（动态获取）
    expect_func_list = expect_func.split(":")
    if len(expect_func_list) < 2:
        # 没有预期函数
        return None
    expect_func_path = expect_func_list[0]
    func_name = expect_func_list[1]
    # 生成模块名（可自定义，如 "my_dynamic_module"）
    module_name = "custom_module"

    # 动态加载模块
    try:
        spec = importlib.util.spec_from_file_location(module_name, expect_func_path)
        custom_module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(custom_module)  # 执行模块代码
    except:
        print(f"load {expect_func_path} fail, {case_desc.case_name} use default golden bin")
        return None

    try:
        custom_func = getattr(custom_module, func_name)
    except AttributeError:
        print(f"load {expect_func_path} fail, {case_desc.case_name} use default golden bin")
        return None
    
    input_list =[]
    output_list = []
    attr_list = []
    for input_desc in case_desc.input_desc_list:
        input_path = input_desc.get("data_path")
        name = input_desc.get("name")

        if input_path is None or len(input_path) == 0:
            input_path = "op_test/{}_{}_{}/input/{}.bin".format(case_desc.op_name.lower() ,case_desc.case_name.lower(), timestamp,name)
        
        if case_desc.case_path and len(case_desc.case_path.strip()) > 0:
            input_path = os.path.join(case_desc.case_path, input_path)

        input_type = input_desc.get("data_type")
        if input_type == "float":
            input_type = "float32"
        elif input_type == "bf16":
            input_type = "bfloat16"

        input_data = np.fromfile(input_path, input_type)
        try:
            input_data = input_data.reshape(input_desc.get("shape"))
        except Exception as e:
            print(f"reshape {input_path} fail, Please verify whether the shape of the input file is consistent with the shape specified in the test case file.")
            return None
        input_list.append(input_data)
    if case_desc.attr_list is None:
        case_desc.attr_list = []

    for attr in case_desc.attr_list:
        attr_list.append(attr.get("value"))
    if len(attr_list) ==0 and len(input_list) >0:
        output_list = custom_func(*input_list)
    elif len(attr_list) >0 and len(input_list) ==0:
        output_list = custom_func(*attr_list)
    elif len(attr_list) >0 and len(input_list) >0:
        output_list = custom_func(*input_list, *attr_list)
    else:
        print(f"case: {case_desc.case_name} input and output is empty, use default golden bin")
        return None
    return output_list
